import warnings

import torch
from transformers import DataCollatorForLanguageModeling
from datasets import load_dataset

warnings.filterwarnings("ignore")
from transformers import AutoModelForMaskedLM
from transformers import AutoTokenizer

model = AutoModelForMaskedLM.from_pretrained('distillbert-base-uncased')


def test_show_param_num():
    distilbert_num_parameters = model.num_parameters() / 1_000_000
    print(f"'>>> DistilBERT number of parameters: {round(distilbert_num_parameters)}M'")
    print(f"'>>> BERT number of parameters: 110M'")


def test_mask():
    text = "This is a great [MASK]."
    tokenizers = AutoTokenizer.from_pretrained("distilbert-base-uncased")
    tokens = tokenizers(text, return_tensors="pt")
    print(tokens)
    print(tokenizers.mask_token_id)
    print(model)
    token_logits = model(**tokens).logits
    print(token_logits.shape)
    # 找到MASK的位置然后获取其预测值
    mask_token_index = torch.where(tokens["input_ids"] == tokenizers.mask_token_id)[1]
    mask_token_logits = token_logits[0, mask_token_index, :]
    # 对MASK所在位置找到他的TOP5预测结果
    top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()
    for token in top_5_tokens:
        print(f"'>>> {text.replace(tokenizers.mask_token, tokenizers.decode([token]))}'")


def test_yingping_data_to_model():
    imdb_dataset = load_dataset("imdb")
    print(imdb_dataset)  # 情感分类数据集，我们只用其中的文本就行了
    sample = imdb_dataset["train"].shuffle(seed=42).select(range(3))
    for row in sample:
        print(f"\n'>>> Review: {row['text']}'")
        print(f"'>>> Label: {row['label']}'")
    # handle data
    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

    def tokenize_function(examples):
        result = tokenizer(examples["text"])
        if tokenizer.is_fast:
            result["word_ids"] = [result.word_ids(i) for i in range(len(result["input_ids"]))]
        return result

    tokenized_datasets = imdb_dataset.map(
        tokenize_function, batched=True, remove_columns=["text", "label"]  # 咱们是完形填空，不需要标签
    )
    print(tokenized_datasets)

    print(tokenizer.model_max_length)  # 预训练模型能设置的最大值 512
    # 我们要做的是完形填空， 所有数据都是宝贵的，我们先拼接数据，在滑动窗口
    chunk_size = 128
    # 看看每一个都多长
    tokenized_samples = tokenized_datasets["train"][:3]
    for idx, sample in enumerate(tokenized_samples["input_ids"]):
        print(f"'>>> Review {idx} length: {len(sample)}'")
    concatenated_examples = {
        k: sum(tokenized_samples[k], []) for k in tokenized_samples.keys()  # 计算拼一起有多少个，
    }
    total_length = len(concatenated_examples["input_ids"])
    print(f"'>>> Concatenated reviews length: {total_length}'")  # 三行的总长度
    chunks = {
        k: [t[i: i + chunk_size] for i in range(0, total_length, chunk_size)]  # 按照咱们刚才指定的chunk_size来切分
        for k, t in concatenated_examples.items()
    }

    for chunk in chunks["input_ids"]:
        print(f"'>>> Chunk length: {len(chunk)}'")


def group_texts(examples, chunk_size=128):
    """处理数据"""
    # 拼接到一起
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
    # 计算长度
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    # //就相当于咱们不要最后多余的了
    total_length = (total_length // chunk_size) * chunk_size
    # 切分
    result = {
        k: [t[i: i + chunk_size] for i in range(0, total_length, chunk_size)]
        for k, t in concatenated_examples.items()
    }
    # 完型填空会用到标签的，也就是原文是啥
    result["labels"] = result["input_ids"].copy()
    return result


def test_handle_data():
    imdb_dataset = load_dataset("imdb")
    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

    def tokenize_function(examples):
        result = tokenizer(examples["text"])
        if tokenizer.is_fast:
            result["word_ids"] = [result.word_ids(i) for i in range(len(result["input_ids"]))]
        return result

    tokenized_datasets = imdb_dataset.map(
        tokenize_function, batched=True, remove_columns=["text", "label"]  # 咱们是完形填空，不需要标签
    )
    lm_datasets = tokenized_datasets.map(group_texts, batched=True)
    print(lm_datasets)
    print(tokenizer.decode(lm_datasets["train"][1]["input_ids"]))
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)  # 0.15是BERT人家说的，咱们别改了
    samples = [lm_datasets["train"][i] for i in range(2)]
    for sample in samples:
        _ = sample.pop("word_ids")  # 不需要wordid
        print(sample)
    for chunk in data_collator(samples)["input_ids"]:
        print(f"\n'>>> {tokenizer.decode(chunk)}'")
        print(len(chunk))


def test_downsample_data():
    train_size = 10000
    test_size = int(0.1 * train_size)
    imdb_dataset = load_dataset("imdb")
    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

    def tokenize_function(examples):
        result = tokenizer(examples["text"])
        if tokenizer.is_fast:
            result["word_ids"] = [result.word_ids(i) for i in range(len(result["input_ids"]))]
        return result

    tokenized_datasets = imdb_dataset.map(
        tokenize_function, batched=True, remove_columns=["text", "label"]  # 咱们是完形填空，不需要标签
    )
    lm_datasets = tokenized_datasets.map(group_texts, batched=True)
    downsampled_dataset = lm_datasets["train"].train_test_split(
        train_size=train_size, test_size=test_size, seed=42
    )
    print(downsampled_dataset)
    return downsampled_dataset


def test_train():
    from transformers import TrainingArguments
    downsampled_dataset = test_downsample_data()
    batch_size = 40
    # 每一个epoch打印结果
    logging_steps = len(downsampled_dataset["train"]) // batch_size
    model_name = "distilbert-base-uncased"

    training_args = TrainingArguments(
        output_dir=f"{model_name}-finetuned-imdb",  # 自己定名字
        overwrite_output_dir=True,
        evaluation_strategy="epoch",
        learning_rate=2e-5,
        weight_decay=0.01,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        logging_steps=logging_steps,
        num_train_epochs=1,
        save_strategy='epoch',
    )
    from transformers import Trainer
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=downsampled_dataset["train"],
        eval_dataset=downsampled_dataset["test"],
        data_collator=data_collator,
    )
    trainer.train()
    # 评估
    import math
    eval_results = trainer.evaluate()
    print(f">>> Perplexity: {math.exp(eval_results['eval_loss']):.2f}")  # 困惑都就是交叉熵的指数形式
    # 这东西有点难理解，用我的话就是你不得在mask那挑啥词合适吗，平均挑了多少个才能答对


def test_model():
    model_checkpoint = "distilbert-base-uncased"
    model = AutoModelForMaskedLM.from_pretrained("./distilbert-base-uncased-finetuned-imdb/checkpoint-84")
    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
    text = 'this is a great [MASK]'
    inputs = tokenizer(text, return_tensors="pt")
    token_logits = model(**inputs).logits
    mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]
    mask_token_logits = token_logits[0, mask_token_index, :]
    top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()
    for token in top_5_tokens:
        print(f"'>>> {text.replace(tokenizer.mask_token, tokenizer.decode([token]))}'")


if __name__ == '__main__':
    # test_show_param_num()
    # test_mask()
    # test_yingping_data_to_model()
    test_downsample_data()
    # test_model()
    # test_train()
